371D - Vessels - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;

#define MAX (int)2e5 + 10

int n;
int capacity[MAX], root[MAX], litersIn[MAX];

int findRoot(int x) {
    if(root[x] == x) {
        return x;
    }
    
    return root[x] = findRoot(root[x]);
}

void _union(int u, int v) {
    int pu = findRoot(u);
    int pv = findRoot(v);
    
    if(pu != pv) {
        root[pu] = pv;
    }
}

void addWater(int vessel, int litersToAdd) {
    if(litersToAdd == 0 || vessel == n+1)
        return;
    
    int canBeAdded = min(capacity[vessel] - litersIn[vessel], litersToAdd);
    
    litersIn[vessel] += canBeAdded;
    
    if(litersIn[vessel] == capacity[vessel]) {
        _union(vessel, vessel+1);
    }
    
    
    litersToAdd -= canBeAdded;
    
    if(!litersToAdd) {
        return;
    }
    
    addWater(findRoot(vessel), litersToAdd);
}

void reset() {
    for(int i=1; i<=n+1; i++) {
        root[i] = i;
        litersIn[i] = 0;
    }
}

int main() {
#ifndef ONLINE_JUDGE
    freopen("/Users/mohtasin.m/Documents/Excercises/SWC/input.txt", "r", stdin);
    freopen("/Users/mohtasin.m/Documents/Excercises/SWC/output.txt", "w", stdout);
#endif
    
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    
    while(cin>>n) {
        reset();
        
        for(int i=1; i<=n; i++) {
            cin>>capacity[i];
        }
        
        int k;
        cin>>k;
        
        while(k--) {
            int query, vessel;
            
            cin>>query>>vessel;
            
            if(query == 2) {
                cout<<litersIn[vessel]<<endl;
                continue;
            }
            
            int litersToAdd;
            cin>>litersToAdd;
            
            addWater(vessel, litersToAdd);
        }
    }
}